#ifndef _contest_prime_h_
#define _contest_prime_h_

#include<vector>
#include<contest/pow.h>

using namespace std;

/**
 * Miller Rabin primality test
 * @see: http://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test 
 **/
template <class I> bool is_probably_prime(const I n, const I* a, const size_t length)
{
    if (n < 2)
        return false;
    for (size_t i=0; i<length; ++i)
        if (n == a[i])
            return true;
    if (n % 2 == 0)
        return false;
    I d = n - 1;
    size_t s;
    for (s = 0; !(d & 1); ++s)
        d >>= 1;
    for (size_t i=0; i<length; ++i)
    {
        I x = pow(a[i], d, n);
        if (x == 1 || x == n - 1)
            continue;
        for (size_t r=0; r<s-1; ++r)
        {
            //Use pow rather than x * x to avoid overflow
            x = pow(x, 2, n);
            if (x == 1)
                return false;
            if (x == n - 1)
                break;
        }
        if (x == n - 1)
            continue;
        return false;
    }
    return true;
}

bool is_prime(const int32_t n)
{
    static const int32_t a[] = {2, 7, 61};
    return is_probably_prime(n, &a[0], 3);
}

bool is_prime(const uint32_t n)
{
    static const uint32_t a[] = {2, 7, 61};
    return is_probably_prime(n, &a[0], 3);
}

bool is_prime(const int64_t n)
{
    static const int64_t a[] = {2, 3, 5, 7, 11, 13, 17};
    return is_probably_prime(n, &a[0], 7);
}

#endif
